from easy_io import read_pkl_file, H5Writer, write_pkl_file
import numpy as np
import h5py
from scipy.ndimage import zoom


def read_partly_from_h5_file(h5_file_obj):
    def helper(path_starts_shape):
        path, starts, shape = path_starts_shape
        slices = tuple(slice(max(st, 0), st + sp) for st, sp in zip(starts, shape))
        raw = h5_file_obj[path][slices]
        mshape = h5_file_obj[path].shape
        pad_width = np.asarray([(0 - st, st + sp - msp) for st, sp, msp in zip(starts, shape, mshape)], 'int')
        pad_width[pad_width < 0] = 0
        paded = np.pad(raw, pad_width, 'constant', constant_values=-1024)
        assert paded.shape == tuple(shape)
        return paded

    return helper


def gen(candidate_pkl_file, data_h5_files, min_edge_length):
    candidates = read_pkl_file(candidate_pkl_file)
    data_file_objs = dict((k, h5py.File(v, 'r')) for k, v in data_h5_files.items())
    readers = dict((k, read_partly_from_h5_file(v)) for k, v in data_file_objs.items())

    for c in candidates:
        source = c['source']
        f = data_file_objs[source]
        reader = readers[source]

        path = c['path']
        center = np.asarray(c['center'])
        diameter = np.asarray(c['diameter'])

        scanid = c['scanid']
        spacing = np.asarray(f[scanid].attrs['spacing'])
        assert spacing[0] == spacing[1], scanid

        unit = np.min(spacing)
        edge_length = max(min_edge_length, diameter)
        crop_shape = np.asarray(np.ceil((edge_length - 1) * unit / spacing + 1), 'int')
        start = np.asarray(np.floor(center - crop_shape / 2 + 1), 'int')
        vol = reader((scanid, start, crop_shape))

        vol = zoom(vol, min_edge_length / np.asarray(vol.shape), output='float32', order=1, mode='nearest')
        assert np.all(np.asarray(vol.shape) == min_edge_length)
        yield path, vol

    for f in data_file_objs.values():
        f.close()


if __name__ == '__main__':
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v7.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v7.pkl',
    #         data_h5_files={
    #             'lidc': '/data_4t/Kaggle/backup/lidc/vol.hdf5',
    #             'kaggle': '/data_4t/Kaggle/backup/kaggle/vol.hdf5'
    #         },
    #         min_edge_length=68,
    #     )
    # )
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/spie_candidates_v2.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/spie_candidates_v2.pkl',
    #         data_h5_files={
    #             'spie': '/data_4t/Kaggle/backup/spie/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v2.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v2.pkl',
    #         data_h5_files={
    #             'kaggle': '/data_4t/Kaggle/backup/kaggle/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v2.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v2.pkl',
    #         data_h5_files={
    #             'kaggle': '/data_4t/Kaggle/backup/kaggle/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )

    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.pkl',
    #         data_h5_files={
    #             'lidc': '/data_4t/Kaggle/backup/lidc/vol.hdf5',
    #             'kaggle_mali': '/data_4t/Kaggle/backup/kaggle/vol.hdf5'
    #         },
    #         min_edge_length=68,
    #     )
    # )
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.hdf5',
    #     'a',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/spie_candidates_v4.pkl',
    #         data_h5_files={
    #             'spie': '/data_4t/Kaggle/backup/spie/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.hdf5',
    #     'a',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v4.pkl',
    #         data_h5_files={
    #             'kaggle_testset': '/data_4t/Kaggle/backup/kaggle/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.hdf5',
    #     'a',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v4.pkl',
    #         data_h5_files={
    #             'kaggle_beni': '/data_4t/Kaggle/backup/kaggle/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )

    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9_nomerge.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9_nomerge.pkl',
    #         data_h5_files={
    #             'lidc': '/data_4t/Kaggle/backup/lidc/vol.hdf5',
    #             'kaggle_mali': '/data_4t/Kaggle/backup/kaggle/vol.hdf5'
    #         },
    #         min_edge_length=68,
    #     )
    # )

    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v11.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v11.pkl',
    #         data_h5_files={
    #             'lidc': '/data_4t/Kaggle/backup/lidc/vol.hdf5',
    #             'kaggle_mali': '/data_4t/Kaggle/backup/kaggle/vol.hdf5'
    #         },
    #         min_edge_length=68,
    #     )
    # )

    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13.pkl',
    #         data_h5_files={
    #             'lidc': '/data_4t/Kaggle/backup/lidc/vol.hdf5',
    #             'kaggle_mali': '/data_4t/Kaggle/backup/kaggle/vol.hdf5'
    #         },
    #         min_edge_length=68,
    #     )
    # )

    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13_nomerge.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13_nomerge.pkl',
    #         data_h5_files={
    #             'lidc': '/data_4t/Kaggle/backup/lidc/vol.hdf5',
    #             'kaggle_mali': '/data_4t/Kaggle/backup/kaggle/vol.hdf5'
    #         },
    #         min_edge_length=68,
    #     )
    # )

    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13_nomerge.hdf5',
    #     'a',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v8_nomerge.pkl',
    #         data_h5_files={
    #             'kaggle_beni': '/data_4t/Kaggle/backup/kaggle/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )
    #
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13_nomerge.hdf5',
    #     'a',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v8_nomerge.pkl',
    #         data_h5_files={
    #             'kaggle_testset': '/data_4t/Kaggle/backup/kaggle/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )
    #
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v13_nomerge.hdf5',
    #     'a',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/spie_candidates_v8_nomerge.pkl',
    #         data_h5_files={
    #             'spie': '/data_4t/Kaggle/backup/spie/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )

    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v15.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v15.pkl',
    #         data_h5_files={
    #             'lidc': '/data_4t/Kaggle/backup/lidc/vol.hdf5',
    #             'kaggle_mali': '/data_4t/Kaggle/backup/kaggle/vol.hdf5'
    #         },
    #         min_edge_length=68,
    #     )
    # )

    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v15.hdf5',
    #     'a',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v10.pkl',
    #         data_h5_files={
    #             'kaggle_beni': '/data_4t/Kaggle/backup/kaggle/vol.hdf5',
    #         },
    #         min_edge_length=68,
    #     )
    # )

    H5Writer(
        '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v15.hdf5',
        'a',
        gen(
            candidate_pkl_file='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v10.pkl',
            data_h5_files={
                'kaggle_testset': '/data_4t/Kaggle/backup/kaggle/vol.hdf5',
            },
            min_edge_length=68,
        )
    )

    H5Writer(
        '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v15.hdf5',
        'a',
        gen(
            candidate_pkl_file='/ssd_1t/huzq/kaggle_data/spie_candidates_v10.pkl',
            data_h5_files={
                'spie': '/data_4t/Kaggle/backup/spie/vol.hdf5',
            },
            min_edge_length=68,
        )
    )
